# Load pckgs for this R session

# --- General 
library(here) # A Simpler Way to Find Your Files   
library(skimr)    # Compact and Flexible Summaries of Data
library(paint)   # Fancy structure for data frames 
library(janitor)  # Simple Tools for Examining and Cleaning Dirty Data
library(tidyverse) # Easily Install and Load the 'Tidyverse'
library(gt) # Easily Create Presentation-Ready Display Tables
library(marginaleffects) # Marginal Effects for Model Objects
library(WeightIt) # Weighting for Covariate Balance in Observational Studies
# --- Plotting & data visualization
library(ggfortify)     # Data Visualization Tools for Statistical Analysis Results
library(ggpubr)        # 'ggplot2' Based Publication Ready Plots
library(patchwork) # The Composer of Plots
# --- Descriptive statistics and regressions
library(broom) # Convert Statistical Objects into Tidy Tibbles
library(Hmisc) # Harrell Miscellaneous
library(estimatr) # Fast Estimators for Design-Based Inference
library(modelsummary) # Summary Tables and Plots for Statistical Models and Data:
library(kableExtra) # Construct Complex Table with 'kable' and Pipe Syntax
library(sandwich) #for robust variance estimation
library(survey) # Analysis of Complex Survey Samples
library(haven) # Import and Export 'SPSS', 'Stata' and 'SAS' Files
library(simstudy) # Simulation of Study Data
library(NHANES) # Data from the US National Health and Nutrition Examination Study  

set.seed(123) # Set seed for reproducibility
data(NHANES)
# Select relevant columns and drop rows with missing values
nhanes_data <- NHANES %>%
  select(ID, Gender, Age, Race1, Height, Weight, BMI, SmokeNow, PhysActive, PhysActiveDays,
         BPSysAve, BPSysAve, BPDiaAve, TotChol, DirectChol, Diabetes, HealthGen,
         Education, HHIncomeMid, Poverty) %>%
  drop_na() %>% 
  # make it smaller 
  slice_sample(n = 1000) # Randomly select 1000 rows using

# Take a quick look at the data
#paint(nhanes_data)

# df DAG 
library(ggdag)
dag_collid <- dagify( Y ~ X + e, Z ~ X + Y, 
                      exposure = "X",outcome = "Y", #latent = "e",
                      # labels
                      labels = c("Z" = "Collider",  "X" = "Exposure",
                                 "Y" = "Outcome", "e" = "Unobserved \nerror"),
                      # positions
                      coords = list(
                        x = c(X = 0, Y= 4, Z = 2 , e = 5),
                        y = c(X = 0, Y = 0, Z = 2, e = 1) 
                      )) %>% 
  tidy_dagitty() 

# Plot DAG 
dag_col_p <-  dag_collid %>% 
  ggdag(., layout = "circle") +  # decided in dagify
  theme_dag_blank(plot.caption = element_text(hjust = 1)) +
  # Nodes 
  geom_dag_node(aes(color = name)) +  
  # Labels dodged to avoid overlap
  geom_dag_label(aes(label = label), color = "#4c4c4c", nudge_x = 0.7, nudge_y = 0.2) +  
  
  geom_dag_text(color="white") +
  labs(title = "Causal map with COLLIDER (Z)" , 
       #subtitle = "X = exposure, Y = outcome, Z = collider, e = random error", 
       caption = ) +
  # Map colors to specific nodes
  scale_color_manual(values = c("X" = "#9b6723", "Y" = "#72aed8", "Z" = "#d02e4c",
                                "e" = "#A6A6A6"), 
                     guide = "none")    # Remove legend

dag_col_p

# df DAG 
dag_confounder <- dagify(Y ~ X + Z + e, X ~  Z, 
                         exposure = "X", outcome = "Y", #latent = "e",
                         # labels
                         labels = c(
                           "Z" = "Confounder",
                           "X" = "Exposure",
                           "Y" = "Outcome",
                           "e" = "Unobserved \nerror"),
                         # positions
                         coords = list(
                           x = c(X = 0, Y= 4, Z = 2 , e = 5),
                           y = c(X = 0, Y = 0, Z = 2, e = 1) 
                         )) %>% 
  tidy_dagitty()     # Add edge_type within pipe

# Plot DAG 
dag_conf_p <-  dag_confounder %>% 
  ggdag(., layout = "circle") + 
  theme_dag_blank(plot.caption = element_text(hjust = 1)) +
  # Nodes 
  geom_dag_node(aes(color = name)) +  
  # Labels dodged to avoid overlap
  geom_dag_label(aes(label = label), color = "#4c4c4c", nudge_x = 0.7, nudge_y = 0.2) +  
  geom_dag_text(color="white") +
  labs(title = "Causal map with CONFOUNDER (Z)" , #subtitle = " ",
        caption = ) +
  # Map colors to specific nodes
  scale_color_manual(values = c("X" = "#9b6723", "Y" = "#72aed8", "Z" = "#d02e4c",
                                "e" = "#A6A6A6"), 
                     guide = "none")   
dag_conf_p

# df DAG 
dag_med <- dagify( Y ~ Z + e,  Z ~ X, exposure = "X", outcome = "Y",  #latent = "e",
                   # labels
                   labels = c(
                     "Z" = "Mediator",
                     "X" = "Exposure",
                     "Y" = "Outcome",
                     "e" = "Unobserved \nerror"),
                   # positions
                   coords = list(
                     x = c(X = 0, Y= 4, Z = 2 , e = 5),
                     y = c(X = 0, Y = 0, Z = 2, e = 1) 
                   )) %>% 
  tidy_dagitty()     # Add edge_type within pipe

# Plot DAG 
dag_med_p <-  dag_med %>% 
  ggdag(., layout = "circle") +  # decided in dagify
  theme_dag_blank(plot.caption = element_text(hjust = 1)) +
  # Nodes 
  geom_dag_node(aes(color = name)) +  
  # Labels dodged to avoid overlap
  geom_dag_label(aes(label = label), color = "#4c4c4c", nudge_x = 0.7, nudge_y = 0.2) +  
  
  geom_dag_text(color="white") +
  labs(title = "Causal map with MEDIATOR (Z)" , 
       #subtitle = "X = exposure, Y = outcome, Z = collider, e = random error", 
       caption = ) +
  # Map colors to specific nodes
  scale_color_manual(values = c("X" = "#9b6723", "Y" = "#72aed8", "Z" = "#d02e4c",
                                "e" = "#A6A6A6"), 
                     guide = "none")    # Remove legend

dag_med_p

# Add predicted values from the unadjusted and adjusted models
nhanes_conf <- nhanes_conf %>%
  mutate(
    pred_unadjusted = predict(model_unadjusted),  # Predicted BPSysAve from unadjusted model
    pred_adjusted = predict(model_adjusted)       # Predicted BPSysAve from adjusted model
  )

# Reshape the data into long format using pivot_longer()
nhanes_long <- nhanes_conf %>%
  pivot_longer(cols = c(pred_unadjusted, pred_adjusted), 
               names_to = "model", values_to = "predicted") %>%
  mutate(model = recode(model, 
                        pred_unadjusted = "Unadjusted: SmokeNow on BPSysAve", 
                        pred_adjusted = "Adjusted: SmokeNow + Age on BPSysAve"))

# Violin plot with dashed line connecting the two conditions
ggplot(nhanes_long, aes(x = factor(SmokeNow), y = BPSysAve, fill = model)) +
  # Create the violin plot to show the distribution of blood pressure
  geom_point(color = "grey", alpha = 0.4, position = position_dodge(width = 0.3) ) +
  # Overlay the predicted dashed line between the two smoking conditions
  geom_line(aes(y = predicted, group = model, color = model), 
            size = 0.8, linetype = "dashed", position = position_dodge(width = 0.3)) +
  # Facet vertically for unadjusted and adjusted models
  facet_wrap(model ~ ., scales = "free_y",ncol = 2) +
  # Add labels and title
  labs(title = "Confounder Analysis: Predicted values in Unadjusted vs Adjusted Regression models",
       subtitle = "Age = Confounder",
       x = "Smoking Status",
       y = "Systolic Blood Pressure (BPSysAve)",
       color = "Model",
       fill = "Model"
       ) +
  # Customize colors for better contrast
  scale_fill_manual(values = c("Unadjusted: SmokeNow on BPSysAve" ="lightcoral" , 
                               "Adjusted: SmokeNow + Age on BPSysAve" = "lightblue")) +
  scale_color_manual(values = c("Unadjusted: SmokeNow on BPSysAve" = "red", 
                                "Adjusted: SmokeNow + Age on BPSysAve" = "blue")) +
  # Minimal theme for clarity
  theme_minimal()+
  theme(legend.position = "none")

# Unadjusted model: total effect of smoking on blood pressure
total_effect_model <- lm(BPSysAve ~ PhysActiveDays, data = nhanes_med)
summary(total_effect_model)$coefficients

# Mediator model: SmokeNow affects BMI
mediator_model <- lm(BMI ~ PhysActiveDays, data = nhanes_med)
summary(mediator_model)$coefficients

# Outcome model: PhysActiveDays and BMI affect BPSysAve
outcome_model <- lm(BPSysAve ~  PhysActiveDays + BMI, data = nhanes_med)
summary(outcome_model)$coefficients

# TOTAL effect: The effect of PhysActiveDays on BPSysAve without adjusting for BMI
total_effect <- coef(total_effect_model)["PhysActiveDays"]# 0.59
total_effect

# DIRECT effect: The effect of PhysActiveDays on BPSysAve after adjusting for BMI.
direct_effect <- coef(outcome_model)["PhysActiveDays"]# 0.621 
direct_effect

# Indirect effect: The portion of the effect on BPSysAve that is mediated through BMI
indirect_effect <- coef(mediator_model)["PhysActiveDays"] * coef(outcome_model)["BMI"] # [-0.0821 X 0.382]  
indirect_effect

# Proportion mediated: The proportion of the total effect that is mediated through BMI
proportion_mediated <- glue::glue("{round(indirect_effect / total_effect *100, 2)}%")
proportion_mediated #  -0.0531

# import data
data <- read_csv(paste0(data_in,"food-health.csv")) # 164,848 

# drop observations
dat <-  data %>%
  filter(age >= 30 & age <60) %>% # 7930 
  # new variables: 
  ## Fruit and vegetables per day (grams)
  ## Blood pressure (systolic+diastolic)
  mutate(fv=veggies_n_fruits_gr,
         bp=blood_pressure) %>%
  filter(fv<3200) %>%
  drop_na(bp) %>% # 7358
  # Days per week exercising 
  mutate(exerc=case_when(paq655 <=7 ~ paq655,
                         paq650 ==2 ~ 0)) %>% 
  # Potato chips per day, grams
  mutate(pchips=gr_potato_chips) %>% 
# select variables
  select(id = seqn, age, gender, weight, height, race, 
         fv, exerc, pchips, 
         bp_systolic, bp_diastolic, bp, bmi,  
         smoker, total_cholesterol, hdl,
         hh_income_usd, hh_income_percap, ln_hh_income_percap ,income_cat
                  ) %>% 
  drop_na() # 7358
 

# Hmisc::describe(dat$hh_income)
paint(dat)